from distance import get_distance
import osmium as osm
import networkx as nx


class OSMHandler(osm.SimpleHandler):

    def __init__(self):
        super(OSMHandler, self).__init__()
        self.G = nx.Graph()

    def node(self, n):
        lon, lat = str(n.location).split('/')
        self.G.add_node(n.id, pos=(lon, lat))

    def way(self, w):
        for i, n in enumerate(w.nodes):
            if i != len(w.nodes) - 1:
                a, b = n.ref, w.nodes[i + 1].ref
                self.G.add_edge(a, b)

    def relation(self, r):
        pass


def write_to_file(osm_handler, info):
    g1 = {Id: i + 1 for i, Id in enumerate(sorted(osm_handler.G.nodes(), key=int))}
    g2 = {i + 1: Id for i, Id in enumerate(sorted(osm_handler.G.nodes(), key=int))}
    with open('./out/'+info+'.txt', 'w', encoding='utf-8') as f:
        f.write('c '+info+' graph\n')
        f.write(f'p edge {len(osm_handler.G.nodes())} {len(osm_handler.G.edges())}\n')
        for node_id in range(1, len(g1) + 1):
            f.write(
                f'n {node_id} {osm_handler.G.node[g2[node_id]]["pos"][0]} {osm_handler.G.node[g2[node_id]]["pos"][1]}\n')
        for edge in osm_handler.G.edges():
            e1, e2 = edge
            f.write(f'e {g1[e1]} {g1[e2]} {get_distance(osm_handler.G, e1, e2)}\n')
